"""
Base Weighting Strategy for Multi-Task Learning
多任务学习权重策略基类
"""

from abc import ABC, abstractmethod
from typing import Dict, List, Any, Optional
import numpy as np
from ml_lib.core import Tensor


class WeightingStrategy(ABC):
    """多任务学习权重策略基类"""
    
    def __init__(self, task_names: List[str], **kwargs):
        """
        初始化权重策略
        
        Args:
            task_names: 任务名称列表
            **kwargs: 其他参数
        """
        self.task_names = task_names
        self.num_tasks = len(task_names)
        self.task_weights = {name: 1.0 for name in task_names}
        self.step_count = 0
        
    @abstractmethod
    def compute_weights(self, losses: Dict[str, Tensor], 
                       gradients: Optional[Dict[str, List[np.ndarray]]] = None,
                       **kwargs) -> Dict[str, float]:
        """
        计算任务权重
        
        Args:
            losses: 各任务的损失值字典
            gradients: 各任务的梯度字典（可选）
            **kwargs: 其他参数
            
        Returns:
            任务权重字典
        """
        pass
    
    def update_weights(self, losses: Dict[str, Tensor], 
                      gradients: Optional[Dict[str, List[np.ndarray]]] = None,
                      **kwargs) -> Dict[str, float]:
        """
        更新任务权重
        
        Args:
            losses: 各任务的损失值字典
            gradients: 各任务的梯度字典（可选）
            **kwargs: 其他参数
            
        Returns:
            更新后的任务权重字典
        """
        self.step_count += 1
        self.task_weights = self.compute_weights(losses, gradients, **kwargs)
        return self.task_weights
    
    def get_weights(self) -> Dict[str, float]:
        """获取当前任务权重"""
        return self.task_weights.copy()
    
    def reset(self):
        """重置权重策略状态"""
        self.task_weights = {name: 1.0 for name in self.task_names}
        self.step_count = 0
    
    def __repr__(self):
        return f"{self.__class__.__name__}(num_tasks={self.num_tasks}, task_names={self.task_names})" 